import logging
import os
import shutil
import random
import numpy as np
from PIL import Image

def parse_filename(filename):
    parts = filename.split('_')
    x_start, x_end = map(int, parts[0].split('-'))
    y_start, y_end = map(int, parts[1].split('-'))
    z_start, z_end = map(int, parts[2].split('-'))
    return x_start, x_end, y_start, y_end, z_start, z_end


def find_pair(file, files):
    x_start, x_end, y_start, y_end, z_start, z_end = parse_filename(file)
    for other_file in files:
        other_x_start, other_x_end, other_y_start, other_y_end, other_z_start, other_z_end = parse_filename(other_file)
        if (x_start == other_x_start and x_end == other_x_end and
            y_start == other_y_start and y_end == other_y_end and
            z_end == other_z_start):
            return other_file
    return None


def calculate_pearson_correlation(file1, file2, directory):
    img1 = Image.open(os.path.join(directory, file1))
    img2 = Image.open(os.path.join(directory, file2))
    arr1 = np.array(img1).flatten()
    arr2 = np.array(img2).flatten()
    return np.corrcoef(arr1, arr2)[0, 1]

def main():
    input_directory = "/root/data/wafer"  # 图片所在目录
    output_directory = "/root/data/random_sample"  # 指定输出目录
    if not os.path.exists(output_directory):
        os.makedirs(output_directory)

    files = [f for f in os.listdir(input_directory) if os.path.isfile(os.path.join(input_directory, f))]
    selected_pairs = 0
    attempts = 0
    max_attempts = 10000  # 防止无限循环
    logging.basicConfig(level=logging.INFO)

    while selected_pairs < 2000 and attempts < max_attempts:
        attempts += 1
        file = random.choice(files)
        pair_file = find_pair(file, files)
        if pair_file:
            correlation = calculate_pearson_correlation(file, pair_file, input_directory)
            if correlation >= 0.3:
                selected_pairs += 1
                new_dir = os.path.join(output_directory, f"{selected_pairs:03d}")
                os.makedirs(new_dir, exist_ok=True)
                shutil.copy(os.path.join(input_directory, file), new_dir)
                shutil.copy(os.path.join(input_directory, pair_file), new_dir)
                logging.info(f"Pair {selected_pairs} found: {file}, {pair_file} (Correlation: {correlation:.2f})")
            else:
                logging.info(f"Pair discarded due to low correlation: {file}, {pair_file} (Correlation: {correlation:.2f})")
        else:
            logging.info(f"No pair found for file: {file}")

        if attempts % 100 == 0:  # 每100次尝试打印一次进度
            logging.info(f"Progress: {attempts} attempts, {selected_pairs} pairs found")


if __name__ == "__main__":
    main()
